# Nuclear Engine Maya Export
# Code by Rockie X.Lee

# Usage:
#   import maya.cmds as cmds
#   cmds.loadPlugin('NEPATH/NEMayaExport.py')
#   cmds.neMayaExport()

import sys
import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMayaAnim as OpenMayaAnim


class NEJoint():
    def __init__(self):
        self.__handle = 0xFFFF
        self.__parent = 0xFFFF
        self.__name = ""
        self.__parentName = ""
        self.__position = OpenMaya.MVector()
        self.__oritation = OpenMaya.MQuaternion()
        self.__scale = OpenMaya.MVector(1.0, 1.0, 1.0)

    def getHandle(self):
        return self.__handle

    def setParent(self, parent):
        self.__parent = parent

    def getParent(self):
        return self.__parent

    def getName(self):
        return self.__name

    def getParentName(self):
        return self.__parentName

    def getPosition(self):
        return self.__position

    def getOritation(self):
        return self.__oritation

    def getScale(self):
        return self.__scale

    def setData(self, handle, dagPath):
        joint = OpenMayaAnim.MFnIkJoint(dagPath)
        self.__handle = handle
        self.__name = joint.name()
        print self.__name
        if joint.parentCount() > 0:
            parent = joint.parent(0)
            if parent.hasFn(OpenMaya.MFn.kJoint):
                parentJoint = OpenMayaAnim.MFnIkJoint(parent)
                self.__parentName = parentJoint.name()
                # print self.__parentName
        self.__position = joint.getTranslation(OpenMaya.MSpace.kPostTransform)
        # print self.__position.x, self.__position.y, self.__position.z
        joint.getRotation(self.__oritation, OpenMaya.MSpace.kPostTransform)
        # print self.__oritation.x, self.__oritation.y, self.__oritation.z, self.__oritation.w
        scalePtr = OpenMaya.MScriptUtil().asDoublePtr()
        """
        joint.getScale(scalePtr)
        self.__scale.x = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 0)
        self.__scale.y = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 1)
        self.__scale.z = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 2)
        """
        # print self.__scale.x, self.__scale.y, self.__scale.z


class NESubMesh():
    def __init__(self):
        self.__space = OpenMaya.MSpace.kWorld
        self.__name = "<>"
        self.__materialName = "<>"

        self.__vertexArray = OpenMaya.MFloatPointArray()
        self.__normalArray = OpenMaya.MFloatVectorArray()
        self.__tangentArray = OpenMaya.MFloatVectorArray()
        self.__colorArray = OpenMaya.MColorArray()
        self.__texCoordUArray = OpenMaya.MFloatArray()
        self.__texCoordVArray = OpenMaya.MFloatArray()
        self.__jointsWeightArray = []
        self.__jointsIndexArray = []

        self.__vertexTempArray = OpenMaya.MFloatPointArray()
        self.__normalTempArray = OpenMaya.MFloatVectorArray()
        self.__tangentTempArray = OpenMaya.MFloatVectorArray()
        self.__colorTempArray = OpenMaya.MColorArray()
        self.__texCoordUTempArray = OpenMaya.MFloatArray()
        self.__texCoordVTempArray = OpenMaya.MFloatArray()
        self.__jointsWeightTempArray = []
        self.__jointsIndexTempArray = []

        self.__hasVertex = False
        self.__hasNormal = False
        self.__hasTangent = False
        self.__hasColor = False
        self.__hasTexCoord = False
        self.__hasWeight = False

        self.__vertexIndexArray = []
        self.__normalIndexArray = []
        self.__tangentIndexArray = []
        self.__colorIndexArray = []
        self.__texCoordIndexArray = []

    def getName(self):
        return self.__name

    def setData(self, dagPath):
        fnMesh = OpenMaya.MFnMesh(dagPath)
        self.__name = fnMesh.name()

        if fnMesh.numVertices() > 0:
            fnMesh.getPoints(self.__vertexTempArray, self.__space)
            self.__hasVertex = True

        if fnMesh.numNormals() > 0:
            fnMesh.getNormals(self.__normalTempArray, self.__space)
            self.__hasNormal = True

        if fnMesh.numUVs() > 0:
            fnMesh.getTangents(self.__tangentTempArray, self.__space)
            self.__hasTangent = True

        if fnMesh.numColors() > 0:
            fnMesh.getColors(self.__colorTempArray)
            self.__hasColor = True

        if fnMesh.numUVs() > 0:
            fnMesh.getUVs(self.__texCoordUTempArray, self.__texCoordVTempArray)
            self.__hasTexCoord = True

        tangentID = 0;
        iptr = OpenMaya.MScriptUtil().asIntPtr()
        for i in range(fnMesh.numPolygons()):
            numVertex = fnMesh.polygonVertexCount(i)
            vertexIndex = OpenMaya.MIntArray()
            fnMesh.getPolygonVertices(i, vertexIndex)
            self.__vertexIndexArray.append(vertexIndex)

            if self.__hasNormal:
                normalIndex = OpenMaya.MIntArray()
                fnMesh.getFaceNormalIds(i, normalIndex)
                self.__normalIndexArray.append(normalIndex)

            if self.__hasTangent:
                tangentIndex = OpenMaya.MIntArray()
                for j in range(numVertex):
                    tangentIndex.append(tangentID)
                    tangentID += 1
                self.__tangentIndexArray.append(tangentIndex)

            if self.__hasColor:
                colorIndex = OpenMaya.MIntArray()
                for j in range(numVertex):
                    fnMesh.getFaceVertexColorIndex(i, j, iptr)
                    colorIndex.append(OpenMaya.MScriptUtil.getInt(iptr))
                self.__colorIndexArray.append(colorIndex)

            if self.__hasTexCoord:
                texCoordIndex = OpenMaya.MIntArray()
                for j in range(numVertex):
                    fnMesh.getPolygonUVid(i, j, iptr)
                    texCoordIndex.append(OpenMaya.MScriptUtil.getInt(iptr))
                self.__texCoordIndexArray.append(texCoordIndex)

    def setJointWeight(self, dagPath, skinCluster, jointHandleMap):
        numWeight = OpenMaya.MScriptUtil().asUintPtr()
        geomIter = OpenMaya.MItGeometry(dagPath)
        self.__hasWeight = True
        while not geomIter.isDone():
            component = geomIter.component()
            weights = OpenMaya.MFloatArray()
            skinCluster.getWeights(dagPath, component, weights, numWeight)
            joints = OpenMaya.MDagPathArray()
            skinCluster.influenceObjects(joints)
            jointsWeightArray = []
            jointsHandleArray = []
            for i in range(joints.length()):
                jointName = joints[i].partialPathName()
                if jointHandleMap.has_key(jointName) and weights[i] > 0.0:
                    jointsWeightArray.append(weights[i])
                    jointsHandleArray.append(jointHandleMap[jointName])
            self.__jointsWeightTempArray.append(jointsWeightArray)
            self.__jointsIndexTempArray.append(jointsHandleArray)
            geomIter.next()

    def convertData(self):
        indexMap = dict()
        numIndex = 0
        for i in range(len(self.__vertexIndexArray)):
            for j in range(3):
                vertex = self.__vertexTempArray[self.__vertexIndexArray[i][j]]
                index = "V" + str(vertex.x) + "&" + str(vertex.y) + "&" + str(vertex.z)

                if self.__hasNormal:
                    normal = self.__normalTempArray[self.__normalIndexArray[i][j]]
                    index += "N" + str(normal.x) + "&" + str(normal.y) + "&" + str(normal.z)

                if self.__hasTangent:
                    tangent = self.__tangentTempArray[self.__tangentIndexArray[i][j]]
                    index += "T" + str(tangent.x) + "&" + str(tangent.y) + "&" + str(tangent.z)

                if self.__hasColor:
                    color = self.__colorTempArray[self.__colorIndexArray[i][j]]
                    index += "C" + str(color.r) + "&" + str(color.g) + "&" + str(color.b) + "&" + str(color.a)

                if self.__hasTexCoord:
                    texCoordU = self.__texCoordUTempArray[self.__texCoordIndexArray[i][j]]
                    texCoordV = self.__texCoordVTempArray[self.__texCoordIndexArray[i][j]]
                    index += "UV" + str(texCoordU) + "&" + str(texCoordV)

                if self.__hasWeight and self.__vertexTempArray.length() == len(self.__jointsWeightTempArray):
                    jointsWeight = self.__jointsWeightTempArray[self.__vertexIndexArray[i][j]]
                    jointsIndex = self.__jointsIndexTempArray[self.__vertexIndexArray[i][j]]

                if indexMap.has_key(index):
                    self.__vertexIndexArray[i][j] = indexMap[index]
                else:
                    indexMap[index] = numIndex
                    self.__vertexArray.append(vertex)
                    self.__vertexIndexArray[i][j] = numIndex
                    if self.__hasNormal:
                        self.__normalArray.append(normal)
                    if self.__hasTangent:
                        self.__tangentArray.append(tangent)
                    if self.__hasColor:
                        self.__colorArray.append(color)
                    if self.__hasTexCoord:
                        self.__texCoordUArray.append(texCoordU)
                        self.__texCoordVArray.append(texCoordV)
                    if self.__hasWeight and self.__vertexTempArray.length() == len(self.__jointsWeightTempArray):
                        self.__jointsWeightArray.append(jointsWeight)
                        self.__jointsIndexArray.append(jointsIndex)
                    numIndex += 1

    def writeData(self, dir):
        file = open(dir + self.__name + ".submeshpy", "w");
        file.write("Name = " + self.__name + "\n")
        file.write("MaterialName = " + self.__materialName + "\n")
        file.write("NumVertex = " + str(self.__vertexArray.length()) + "\n")
        file.write("NumNormal = " + str(self.__normalArray.length()) + "\n")
        file.write("NumTangent = " + str(self.__tangentArray.length()) + "\n")
        file.write("ColorSize = " + str(16) + "\n")
        file.write("NumColor = " + str(self.__colorArray.length()) + "\n")
        file.write("NumSecondaryColor = " + str(0) + "\n")
        file.write("TexCoordSize = " + str(8) + "\n")
        file.write("NumTexCoord = " + str(self.__texCoordUArray.length()) + "\n")
        file.write("NumBlendCount = " + str(4) + "\n")
        file.write("NumBlendWeight = " + str(len(self.__jointsWeightArray)) + "\n")
        file.write("NumBlendIndex = " + str(len(self.__jointsIndexArray)) + "\n")
        file.write("IndexSize = " + str(4) + "\n")
        file.write("NumIndex = " + str(len(self.__vertexIndexArray) * 3) + "\n")
        file.write("NumFace = " + str(len(self.__vertexIndexArray)) + "\n")

        for i in range(self.__vertexArray.length()):
            vertex = self.__vertexArray[i]
            file.write("V = " + str(vertex.x) + " " + str(vertex.y) + " " + str(vertex.z) + "\n")

        for i in range(self.__normalArray.length()):
            normal = self.__normalArray[i]
            file.write("N = " + str(normal.x) + " " + str(normal.y) + " " + str(normal.z) + "\n")

        for i in range(self.__tangentArray.length()):
            tangent = self.__tangentArray[i]
            file.write("T = " + str(tangent.x) + " " + str(tangent.y) + " " + str(tangent.z) + "\n")

        for i in range(self.__colorArray.length()):
            color = self.__colorArray[i]
            file.write("C = " + str(color.r) + " " + str(color.g) + " " + str(color.b) + " " + str(color.a) + "\n")

        for i in range(self.__texCoordUArray.length()):
            file.write("UV = " + str(self.__texCoordUArray[i]) + " " + str(self.__texCoordVArray[i]) + "\n")

        for i in range(len(self.__jointsWeightArray)):
            file.write("W = ")
            for j in range(len(self.__jointsWeightArray[i])):
                file.write(str(self.__jointsWeightArray[i][j]) + " ")
            if len(self.__jointsWeightArray[i]) < 4:
                for j in range(len(self.__jointsWeightArray[i]), 4):
                    file.write("0.0 ")
            file.write("\n")

        for i in range(len(self.__jointsIndexArray)):
            file.write("J = ")
            for j in range(len(self.__jointsIndexArray[i])):
                file.write(str(self.__jointsIndexArray[i][j]) + " ")
            if len(self.__jointsIndexArray[i]) < 4:
                for j in range(len(self.__jointsIndexArray[i]), 4):
                    file.write("0 ")
            file.write("\n")

        for i in range(len(self.__vertexIndexArray)):
            file.write("F = " + str(self.__vertexIndexArray[i][0]) + " " + str(self.__vertexIndexArray[i][1]) + " " + str(self.__vertexIndexArray[i][2]) + "\n")


neMayaPluginCmdName = "neMayaExport"

neDirFlag = "-d"
neDirLongFlag = "-dir"

# Command
class NEMayaExportCommand(OpenMayaMPx.MPxCommand):
    def __init__(self):
        OpenMayaMPx.MPxCommand.__init__(self)
        self.__dir = ""
        self.__meshName = "Mesh"
        self.__textureName = "<>"
        self.__bumpMapName = "<>"
        self.__subMeshs = []

        self.__skinCluster = 0

        self.__skeletonName = "Skeleton"
        self.__numJoints = 0
        self.__joints = []
        self.__jointHandleMap = dict()

    def doIt(self, argList):
        print "[Begin] Nuclear Engine Maya Export===================================="
        self.export()
        print "[Success] Nuclear Engine Maya Export=================================="

    def export(self):
        # Create a selection list and iterator
        selectList = OpenMaya.MSelectionList()
        OpenMaya.MGlobal.getActiveSelectionList(selectList)
        iter = OpenMaya.MItSelectionList(selectList)
        dagIter = OpenMaya.MItDag(OpenMaya.MItDag.kDepthFirst, OpenMaya.MFn.kInvalid)

        dependNode = OpenMaya.MObject()
        while not iter.isDone():
            iter.getDependNode(dependNode)
            objectPath = OpenMaya.MDagPath()
            iter.getDagPath(objectPath)
            dagIter.reset(objectPath.node(), OpenMaya.MItDag.kDepthFirst, OpenMaya.MFn.kInvalid)
            while not dagIter.isDone():
                dagPath = OpenMaya.MDagPath()
                dagIter.getPath(dagPath)

                if dagPath.hasFn(OpenMaya.MFn.kMesh) and dagPath.hasFn(OpenMaya.MFn.kTransform):
                    pass

                elif dagPath.hasFn(OpenMaya.MFn.kMesh):
                    self.getSkinCluster(dagPath)
                    subMesh = NESubMesh()
                    subMesh.setData(dagPath)
                    if self.__skinCluster != 0:
                        self.__subMeshs.append([subMesh, dagPath, self.__skinCluster])
                    else:
                        self.__subMeshs.append([subMesh, dagPath, 0])

                elif dagPath.hasFn(OpenMaya.MFn.kJoint):
                    joint = NEJoint()
                    joint.setData(self.__numJoints, dagPath)
                    self.__joints.append(joint)
                    self.__jointHandleMap[joint.getName()] = self.__numJoints
                    self.__numJoints += 1

                dagIter.next()
            iter.next()

        self.writeMeshData(self.__dir)
        if self.__numJoints > 0:
            self.setJointsParentHandle()
            self.writeSkeletonData(self.__dir)
        else:
            print "No seleted joints !"

        for it in self.__subMeshs:
            subMesh = it[0]
            dagPath = it[1]
            skinCluster = it[2]
            if self.__numJoints > 0 and skinCluster != 0:
                subMesh.setJointWeight(dagPath, skinCluster, self.__jointHandleMap)
            subMesh.convertData()
            subMesh.writeData(self.__dir)

    def writeMeshData(self, dir):
        file = open(dir + self.__meshName + ".meshpy", "w");
        file.write("Name = " + self.__meshName + "\n")
        file.write("TextureName = " + self.__textureName + "\n")
        file.write("BumpMapName = " + self.__bumpMapName + "\n")
        file.write("NumSubMeshs = " + str(len(self.__subMeshs)) + "\n")
        for i in range(len(self.__subMeshs)):
            file.write("SM = " + self.__subMeshs[i][0].getName() + "\n")
        file.close()
        print "Finished export mesh data !"

    def getSkinCluster(self, dagPath):
        fnMesh = OpenMaya.MFnMesh(dagPath)
        depNodeIt = OpenMaya.MItDependencyNodes(OpenMaya.MFn.kSkinClusterFilter)
        while not depNodeIt.isDone():
            object = depNodeIt.item()
            self.__skinCluster = OpenMayaAnim.MFnSkinCluster(object)
            numGeometries = self.__skinCluster.numOutputConnections()
            for i in range(numGeometries):
                index = self.__skinCluster.indexForOutputConnection(i)
                outObject = self.__skinCluster.outputShapeAtIndex(index)
                if outObject == fnMesh.object():
                    print "Found skin cluster !"
                    return
                else:
                    self.__skinCluster = 0
            depNodeIt.next()

    def setJointsParentHandle(self):
        for joint in self.__joints:
            parentName = joint.getParentName()
            if parentName != "":
                if self.__jointHandleMap.has_key(parentName):
                    joint.setParent(self.__jointHandleMap[parentName])

    def writeSkeletonData(self, dir):
        file = open(dir + self.__skeletonName + ".skeletonpy", "w");
        file.write("Name = " + self.__skeletonName + "\n")
        file.write("NumJoints = " + str(self.__numJoints) + "\n")
        for joint in self.__joints:
            file.write("Handle = " + str(joint.getHandle()) + "\n")
            file.write("Parent = " + str(joint.getParent()) + "\n")
            file.write("Name = " + joint.getName() + "\n")
            position = joint.getPosition()
            file.write("Position = " + str(position.x) + " " + str(position.y) + " " + str(position.z) + "\n")
            oritation = joint.getOritation()
            file.write("Oritation = " + str(oritation.x) + " " + str(oritation.y) + " " + str(oritation.z) + " " + str(oritation.w) + "\n")
            scale = joint.getScale()
            file.write("Scale = " + str(scale.x) + " " + str(scale.y) + " " + str(scale.z) + "\n")
        file.close()
        print "Finished export skeleton data !"

# Creator
def cmdCreator():
    return OpenMayaMPx.asMPxPtr(NEMayaExportCommand())


# Syntax creator
def syntaxCreator():
    syntax = OpenMaya.MSyntax()
    #syntax.addFlag(neDirFlag, neDirLongFlag, OpenMaya.MSyntax.kString)
    return syntax


# Initialize the plug-in
def initializePlugin(mayaObject):
    mayaPlugin = OpenMayaMPx.MFnPlugin(mayaObject)
    try:
        mayaPlugin.registerCommand(neMayaPluginCmdName, cmdCreator, syntaxCreator)
    except:
        sys.stderr.write("Failed to register command: %s\n" % neMayaPluginCmdName)
        raise


# Uninitialize the plug-in
def uninitializePlugin(mayaObject):
    mayaPlugin = OpenMayaMPx.MFnPlugin(mayaObject)
    try:
        mayaPlugin.deregisterCommand(neMayaPluginCmdName)
    except:
        sys.stderr.write("Failed to unregister command: %s\n" % neMayaPluginCmdName)
        raise

